import threading
import multiprocessing
import dynamic_pvalue
import numpy


class Pvalue:
    """
    Class to calculate p-values
    """
    
    @staticmethod    
    def compute_pvalues(proteins, *args):
        """
        Input:
        proteins - list of data_handler.Proteins objects 
        args - list of other arguments (specific to each type of pvalues)
        
        Output: the pvalue field in the proteins is filled
                (inf is used when we failed to compute the pvalue)
        """        
        pass
    
    
class DynamicPvalue:
    """
    Class to calculate p-values using dynamic programming
    """
    
    @staticmethod                    
    def compute_pvalues(proteins, smooth_pvalue, steps_per_position, \
            use_threads):        
        if use_threads:
            DynamicPvalue.compute_threaded_pvalues(proteins, smooth_pvalue, \
                    steps_per_position)        
        else:
            DynamicPvalue.compute_pval_proteins(proteins, smooth_pvalue, \
                    steps_per_position)
            
            
    @staticmethod             
    def compute_pval_proteins(proteins, smooth_pvalue, steps_per_position):
        for protein in proteins:
            pval = dynamic_pvalue.compute_dynamic_pval(protein.probs, \
                    protein.matched, smooth_pvalue, steps_per_position) 
            protein.pvalue = pval
            
    
    @staticmethod       
    def worker(ps, ds, start, smooth_pvalue, steps_per_position, out_q):  
        out_dict = {}
        i = start
        for (p, d) in zip(ps, ds):
            pval = dynamic_pvalue.compute_dynamic_pval(p, d, smooth_pvalue, \
                    steps_per_position)             
            out_dict[i] = pval
            i += 1
        out_q.put(out_dict)     
        
    
    @staticmethod                            
    def compute_threaded_pvalues(proteins,  smooth_pvalue, steps_per_position):                      
        ncpus = max(multiprocessing.cpu_count() - 1, 1)
        n = len(proteins)
        step = int(numpy.ceil(n / float(ncpus)))
        out_q = multiprocessing.Queue()        
        # start a number of threads
        my_threads = []
        for k in range(ncpus):
            start = k * step
            stop = (k+1) * step
            ps, ds = zip(*[(p.probs, p.matched) for p in proteins[start:stop]])
            t = multiprocessing.Process(target = DynamicPvalue.worker, \
                    args = (ps, ds, start, smooth_pvalue, steps_per_position, \
                    out_q)) 
            my_threads.append(t)
            t.start()       
                       
        # merge all the results 
        result_dict = {}
        for k in range(ncpus):
            result_dict.update(out_q.get())                                               
        # Wait for all worker processes to finish
        for p in my_threads:
            p.join()    
            
        k = 0          
        for index, pval in result_dict.iteritems():
            proteins[index].pvalue = pval
            k += 1
        # TO DO: remove k
        print len(proteins), k
